// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/rows_streaming_window_build.h>
#include <pollux/common/testutil/test_value.h>
#include <pollux/exec/window_function.h>

namespace kumo::pollux::exec {
    namespace {
        bool hasRangeFrame(const std::shared_ptr<const core::WindowNode> &windowNode) {
            for (const auto &function: windowNode->windowFunctions()) {
                if (function.frame.type == core::WindowNode::WindowType::kRange) {
                    return true;
                }
            }
            return false;
        }
    } // namespace

    RowsStreamingWindowBuild::RowsStreamingWindowBuild(
        const std::shared_ptr<const core::WindowNode> &windowNode,
        pollux::memory::MemoryPool *pool,
        const common::SpillConfig *spillConfig,
        tsan_atomic<bool> *nonReclaimableSection)
        : WindowBuild(windowNode, pool, spillConfig, nonReclaimableSection),
          hasRangeFrame_(hasRangeFrame(windowNode)) {
        pollux::common::testutil::TestValue::adjust(
            "kumo::pollux::exec::RowsStreamingWindowBuild::RowsStreamingWindowBuild",
            this);
    }

    bool RowsStreamingWindowBuild::needsInput() {
        // We need input if there is no or only partition.
        return windowPartitions_.size() < 2;
    }

    void RowsStreamingWindowBuild::ensureInputPartition() {
        if (windowPartitions_.empty() || windowPartitions_.back()->complete()) {
            windowPartitions_.emplace_back(std::make_shared<WindowPartition>(
                data_.get(), inversedInputChannels_, sortKeyInfo_));
        }
    }

    void RowsStreamingWindowBuild::addPartitionInputs(bool finished) {
        if (inputRows_.empty()) {
            return;
        }

        ensureInputPartition();
        windowPartitions_.back()->addRows(inputRows_);

        if (finished) {
            windowPartitions_.back()->setComplete();
        }

        inputRows_.clear();
        inputRows_.shrink_to_fit();
    }

    void RowsStreamingWindowBuild::addInput(RowVectorPtr input) {
        for (auto i = 0; i < inputChannels_.size(); ++i) {
            decodedInputVectors_[i].decode(*input->childAt(inputChannels_[i]));
        }

        for (auto row = 0; row < input->size(); ++row) {
            char *newRow = data_->newRow();

            for (auto col = 0; col < input->childrenSize(); ++col) {
                data_->store(decodedInputVectors_[col], row, newRow, col);
            }

            if (previousRow_ != nullptr &&
                compareRowsWithKeys(previousRow_, newRow, partitionKeyInfo_)) {
                addPartitionInputs(true);
            }

            if (previousRow_ != nullptr && inputRows_.size() >= numRowsPerOutput_) {
                // Needs to wait the peer group ready for range frame.
                if (hasRangeFrame_) {
                    if (compareRowsWithKeys(previousRow_, newRow, sortKeyInfo_)) {
                        addPartitionInputs(false);
                    }
                } else {
                    addPartitionInputs(false);
                }
            }

            inputRows_.push_back(newRow);
            previousRow_ = newRow;
        }
    }

    void RowsStreamingWindowBuild::noMoreInput() {
        addPartitionInputs(true);
    }

    std::shared_ptr<WindowPartition> RowsStreamingWindowBuild::nextPartition() {
        // Remove the processed output partition from the queue.
        //
        // NOTE: the window operator only calls this after processing a completed
        // partition.
        if (!windowPartitions_.empty() && windowPartitions_.front()->complete() &&
            windowPartitions_.front()->numRows() == 0) {
            windowPartitions_.pop_front();
        }

        POLLUX_CHECK(hasNextPartition());
        return windowPartitions_.front();
    }

    bool RowsStreamingWindowBuild::hasNextPartition() {
        // Checks if there is a window partition that is either incomplete or
        // completed but has unconsumed rows.
        for (auto it = windowPartitions_.rbegin(); it != windowPartitions_.rend();
             ++it) {
            const auto &windowPartition = *it;
            if (!windowPartition->complete() || windowPartition->numRows() > 0) {
                return true;
            }
        }

        return false;
    }
} // namespace kumo::pollux::exec
